import torchvision.transforms as transforms

from data.mnist import MNIST
from data.cifar import CIFAR10, CIFAR100


def load_dataset(dataset, noise_type, noise_rate):
    # Load dataset
    if dataset == 'mnist':
        transform_train = transforms.Compose([
            transforms.RandomCrop(28, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
        ])
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,)),
        ])
        input_channel = 1
        num_classes = 10
        train_dataset = MNIST(
            root='./data/',
            download=True, 
            train=True, 
            transform=transform_train, 
            noise_type=noise_type,
            noise_rate=noise_rate, 
        )
        test_dataset = MNIST(
            root='./data/',
            download=True, 
            train=False, 
            transform=transform_test, 
        )
        
    if dataset == 'cifar10':
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
        ])
        input_channel = 3
        num_classes = 10
        train_dataset = CIFAR10(
            root='./data/', 
            download=True, 
            train=True, 
            transform=transform_train, 
            noise_type=noise_type, 
            noise_rate=noise_rate, 
        )
        test_dataset = CIFAR10(
            root='./data/',
            download=True, 
            train=False, 
            transform=transform_test, 
        )

    if dataset == 'cifar100':
        transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762)),
        ])
        transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762)),
        ])
        input_channel = 3
        num_classes = 100
        train_dataset = CIFAR100(
            root='./data/', 
            download=True, 
            train=True, 
            transform=transform_train, 
            noise_type=noise_type, 
            noise_rate=noise_rate, 
        )
        test_dataset = CIFAR100(
            root='./data/',
            download=True, 
            train=False, 
            transform=transform_test, 
        )
    
    return input_channel, num_classes, train_dataset, test_dataset
